"""Test CodeHierarchyNodeParser with skeleton option set to False."""

import os
from typing import List, cast

from llama_index.packs.code_hierarchy import CodeHierarchyNodeParser
from llama_index.core.schema import NodeRelationship, RelatedNodeInfo, TextNode
from llama_index.packs.code_hierarchy import CodeHierarchyAgentPack
from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.core import  GPTVectorStoreIndex,VectorStoreIndex
from llama_index.llms import openai_like
from llama_index.core import Settings
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.huggingface import HuggingFaceEmbedding  # HuggingFaceEmbedding:用于将文本转换为词向量
from llama_index.llms.huggingface import HuggingFaceLLM  # HuggingFaceLLM：用于运行Hugging Face的预训练语言模型
from llama_index.core import Settings,SimpleDirectoryReader,VectorStoreIndex
import chromadb
from llama_index.embeddings.dashscope import DashScopeEmbedding
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext, load_index_from_storage
from llama_index.llms.deepseek  import DeepSeek
from llama_index.embeddings.fastembed import FastEmbedEmbedding


llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm
embed_model = FastEmbedEmbedding(model_name="BAAI/bge-small-en-v1.5")
Settings.embed_model = embed_model

def test_python_code_splitter() -> None:
    """Test case for code splitting using python."""
    if "CI" in os.environ:
        return

    code_splitter = CodeHierarchyNodeParser(
        language="python", skeleton=False, chunk_min_characters=0
    )

    text = """\
class Foo:
    def bar() -> None:
        print("bar")

    async def baz():
        print("baz")"""

    text_node = TextNode(
        text=text,
        metadata={
            "module": "example.foo",
        },
    )

    chunks: List[TextNode] = code_splitter.get_nodes_from_documents([text_node])
   
    pack= CodeHierarchyAgentPack(chunks,llm=llm)
    output=pack.run("文档中有那些函数？")
    print("output:"+output)
 

 

test_python_code_splitter()